-- YUVDiff

import Char
import qualified Data.ByteString.Lazy as B
import Data.Word
import Foreign ( unsafePerformIO )
import Text.Printf
import Text.Regex
import List
import Maybe
import System
import System.IO
import Control.Monad
import YUVPicture


ifo = (YUV420,352,288) :: PicInfo
tgt = "D:/Work/_/w0801/HostCodec/decoder/prj/c6416/Debug/trace.yuv" :: FilePath
ref = "D:/Work/_/w0801/HostCodec/decoder/test/stuff/test.yuv" :: FilePath
ori = "D:/Work/_/w0801/HostCodec/decoder/test/stuff/trace_dec.nlf.yuv" :: FilePath

pic1 = (YUV420,[y,u,v]) :: Picture
  where
    y = map B.pack . chop 6 $ [1..24]
    u = map B.pack . chop 3 $ [31..36]
    v = map B.pack . chop 3 $ [51..56]

lastTest = seeDiff ifo tgt ori . fromJust $ findDiff ifo tgt ori


showCmdUsage = unlines [
 	"To locate / show difference pixel between YUV seq files.",
  "",
  "Usage:",
  "\tYUVDiff.exe size-info rec-yuv dec-yuv [F-MX-MY]",
  "",
  "size-info:",
  "\t-spWxH\t\tW*H pixels.",
  "\t-smWxH\t\tW*H macroblocks.",
  "\t-mM\t\tof M mode, where M may be:",
  "\t\t\t  QCIF, CIF, D1 (case-insensitive)",
  "",
  "The last optional argument is to specify which MB to show.",
  "If this is not specified, I will show you the first diff-MB.",
  "Where F, MX and MY are all numbers.",
  "\tF\t\tframe number",
  "\tMX\t\tX coordinate of MB",
  "\tMY\t\tY coordinate of MB",
  "",
  "Version 2.0",
  "Made by Levi.G. 2008-12. All rights reserved.",
  "Email: leaveye <dot> guo <at> gmail <dot> com" ]

showHistory = unlines [
  "Ver 1.0",
  "2007-06",
  "Functional work fine, using getContent function.",
  "Slow and need much more memory.",
  "",
  "Ver 2.0",
  "2008-12",
  "Rewrite using Data.ByteString.Lazy, hGet function.",
  "Much more faster than Ver 1.0." ]

main = getArgs >>= dealArgs
  where
    dealArgs args
      | length args >= 1 && isPrefixOf "-h" (args !! 0)
                        = giveHistory
      | length args < 3 = giveUsage
      | ( isNothing . readOption . head ) args
                        = giveUsage
      | otherwise       = case mbPos of
            Just p  -> seeDiff ifo fn1 fn2 p >> return ()
            Nothing -> putStrLn "There is no conflict data in these two YUV file."
      where
        fn1 = args !! 1
        fn2 = args !! 2
        (w,h) = fromJust . readOption . head $ args
        ifo   = (YUV420,w,h)
        mbPos = if length args > 3
                  then readPossition (args !!3)
                  else findDiff ifo fn1 fn2

    giveUsage = putStr showCmdUsage
    giveHistory = putStr showHistory

    match = matchRegex . mkRegex

    readOption :: String -> Maybe (Int,Int)
    readOption ('-':'s':'p':sz) = case match "([0-9]+)x([0-9]+)" sz of
            Just [ws,hs] -> Just (read ws, read hs)
            _            -> Nothing
    readOption ('-':'s':'m':sz) = case match "([0-9]+)x([0-9]+)" sz of
            Just [ws,hs] -> Just (read ws * 16, read hs * 16)
            _            -> Nothing
    readOption ('-':'m':m:_) = case toUpper m of
            'Q' ->  readOption "-sp176x144"
            'C' ->  readOption "-sp352x288"
            'D' ->  readOption "-sp720x576"
            _   ->  Nothing
    readOption _ = Nothing

    readPossition :: String -> Maybe (Int,(Int,Int))
    readPossition opt = case match "([0-9]+)-([0-9]+)-([0-9]+)" opt of
            Just [f,mx,my] -> Just (read f, (read mx, read my))
            _              -> Nothing

genL f = snd . mapAccumL f (-4::Int)

genStrAcc r str
  | r == 19         = ( r+1, genPad . concat $ genL (genCharAcc '_' ' ' '|' ' ') str )
  | r == -1         = ( r+1, genPad . concat $ genL (genCharAcc '_' '_' '|' ' ') str )
  | r `rem` 4 == 3  = ( r+1, genPad . concat $ genL (genCharAcc '_' '_' '|' '|') str )
  | r < 0 || r > 15 = ( r+1, genPad . concat $ genL (genCharAcc ' ' ' ' '|' ' ') str )
  | otherwise       = ( r+1, genPad . concat $ genL (genCharAcc ' ' ' ' '|' '|') str )
  where
    genCharAcc d1 d2 d3 d4 i c
      | i == 19         = ( i+1, dx:cc:d4:[] )
      | i == -4         = ( i+1, d4:cc:[] )
      | i `rem` 4 == 0  = ( i+1, d3:cc:[] )
      | i < 0 || i > 15 = ( i+1, d2:cc:[] )
      | otherwise       = ( i+1, d1:cc:[] )
      where cc = if c == ' ' then dx else c
            dx = if i < 0 || i > 15 then d2 else d1
    genPad s = printf "%01X%c%s%c%01X" rn dp s dp rn :: String
      where rn = r `mod` 16
            dp = if r `mod` 4 == 3 && r < 19 then '_' else ' '

seeDiff :: PicInfo -> FilePath -> FilePath -> (Int,(Int,Int)) -> IO ([[(Word8,Word8,Word8)]],[[(Word8,Word8,Word8)]])
seeDiff ifo fn1 fn2 (fn,mb@(mx,my)) = do
    p1 <- getPic ifo fn1 fn >>= toPix
    p2 <- getPic ifo fn2 fn >>= toPix
    printf "File LEFT  : \"%s\"\n" fn1
    printf "File RIGHT : \"%s\"\n" fn2
    printf "Position : frame no = %d, MB (x,y) = (%d,%d)\n" fn mx my
    printf "\n%61s\n" "____/[ (diff number) = (YUV) in binary ]\\____"
    printf "%62s\n" "C D E F|0_1_2_3|4_5_6_7|8_9_A_B|C_D_E_F|0 1 2 3"
    let pixDiffArray = zipWith (zipWith getMark) p1 p2
    mapM_ (printf "%65s\n") $ genL genStrAcc pixDiffArray
    printf "%62s\n" "C D E F|0 1 2 3|4 5 6 7|8 9 A B|C D E F|0 1 2 3"
    let pixDiffList = intercalate [showDiffItemLabel] . filter (not . null)
                    . map (map (\(a,Just b) -> showDiffItem mb a b)
                         . filter (isJust . snd))
                    . zipWith (\r l -> flip zip l $ zip [(-4)..] (repeat r)) [(-4)..]
                    $ zipWith (zipWith showDiff) p1 p2
    --putStrLn showDiffItemLabel
    --putStr . unlines $ pixDiffList
    --putStrLn showDiffItemLabel
    return (p1,p2)
  where
    tableHdr = "CDEF0123456789ABCDEF0123"
    pad n s = replicate (n - length s) ' ' ++ s
    intercalate a = concat . intersperse a
    diffItemWidths = [0,4,4,5,5,5,5,5,5,5,0]
    showDiffItemLabel = intercalate " | " . zipWith pad diffItemWidths $
        ["","MB X","MB Y","Off X","Off Y","Pix X","Pix Y","  Y  ","  U  ","  V  ",""]
    showDiffItem mb@(mx,my) off@(ox,oy) dif@[dy,du,dv] = intercalate " | " . zipWith pad diffItemWidths $
        ["", show mx, show my, show ox, show oy, show (mx*16+ox), show (my*16+oy), dy, du, dv, ""]
    showDiff m@(a,b,c) n@(x,y,z)
      | m == n    = Nothing
      | otherwise = Just $ zipWith showCmp [a,b,c] [x,y,z]
      where
        showCmp a b | a /= b    = showHex a ++ '|' : showHex b
                    | otherwise = replicate 3 ' ' ++ showHex a
        showHex x = [toHex (x `div` 16), toHex (x `rem` 16)]
        toHex = fromJust . flip lookup (zip [0..15] "0123456789ABCDEF")
    getMark x1@(y1,u1,v1) x2@(y2,u2,v2)
      | x1 == x2  = ' '
      | otherwise = head . show
                  . sum . map snd . filter fst . flip zip [4,2,1]
                  $ zipWith (/=) [y1,u1,v1] [y2,u2,v2]
    toPix = return
          . map (map (\[y,u,v]->(y,u,v)) . transpose) . transpose
          . map (map B.unpack) . snd . toYUV444 . cropP mb . fromJust
    cropP (mbX,mbY) = cropPic (x0,y0) (x1,y1)
      where x0 = mbX * 16 - 4
            y0 = mbY * 16 - 4
            x1 = mbX * 16 + 20
            y1 = mbY * 16 + 20

findDiff :: PicInfo -> FilePath -> FilePath -> Maybe (Int,(Int,Int))
findDiff c fn1 fn2 = unsafePerformIO $ findDiffIO c fn1 fn2

findDiffIO :: PicInfo -> FilePath -> FilePath -> IO (Maybe (Int,(Int,Int)))
findDiffIO ifo@(m,w,h) fn1 fn2 = do
    fd1 <- openBinaryFile fn1 ReadMode
    fd2 <- openBinaryFile fn2 ReadMode
    dPic <- findDiffPicIO 0 ifo fd1 fd2
    case dPic of
      Just (frameNo,(pic1,pic2)) -> do
        let Just (mb@(mbX,mbY)) = findDiffMB pic1 pic2
        return $ Just (frameNo,mb) -- ( ***, (cropP mb pic1, cropP mb pic2) )
      Nothing -> return Nothing
  where
    (rWCb,rWCr) = scaleWidth m
    (rHCb,rHCr) = scaleHeight m
    scale sz = round . (*) (toRational sz)
    crop (mbX,mbY) (rW,rH) = cropPlane (x0,y0) (x1,y1)
      where x0 = flip scale rW $ mbX * 16 - 4
            y0 = flip scale rH $ mbY * 16 - 4
            x1 = flip scale rW $ mbX * 16 + 20
            y1 = flip scale rH $ mbY * 16 + 20
    cropP (mbX,mbY) = cropPic (x0,y0) (x1,y1)
      where x0 = mbX * 16 - 4
            y0 = mbY * 16 - 4
            x1 = mbX * 16 + 20
            y1 = mbY * 16 + 20

findDiffPicIO :: Int -> PicInfo -> Handle -> Handle -> IO (Maybe (Int,(Picture,Picture)))
findDiffPicIO fn ifo fd1 fd2 =
  if any (unsafePerformIO . hIsEOF) [fd1,fd2]
    then return Nothing
    else do
      p1 <- hGetPic ifo fd1
      p2 <- hGetPic ifo fd2
      if p1 == p2
        then findDiffPicIO (fn+1) ifo fd1 fd2
        else return $ Just (fn,(p1,p2))

findDiffMB :: Picture -> Picture -> Maybe (Int,Int)
findDiffMB p1 p2 = do
    (y,(r1,r2)) <- findNe 0 (unpackPic p1) (unpackPic p2)
    (x,_)       <- findNe 0 r1 r2
    return (x,y)
  where
    mbs1 = unpackPic p1
    mbs2 = unpackPic p2
    findNe i xs ys
      | any null [xs,ys]   = Nothing
      | head xs /= head ys = Just ( i, (head xs, head ys) )
      | otherwise          = findNe (i+1) (tail xs) (tail ys)
